from python_nlp.kg.neo4j.config import NEO4J_CONFIG_CBLUE_DETAIL
from neo4j import GraphDatabase
import json
from PyCmpltrtok.common import sep

if '__main__' == __name__:

    def _load(driver, xpath, limit=0, index=0):
        sep(f'{xpath},limit={limit}')
        with driver.session() as ss:
            with open(xpath, 'r', encoding='utf8') as fi:
                cnt = -1
                while True:
                    # 读取一行
                    xline = fi.readline()
                    if not xline:
                        break
                    if '\r\n' == xline[-2:]:
                        xline = xline[:-2]
                    elif '\r' == xline[-1:]:
                        xline = xline[:-1]
                    xline = xline.strip()
                    if not xline:
                        continue

                    # 人为限定读取的行数
                    cnt += 1
                    if limit and cnt >= limit:
                        break

                    # 让运行者知道处理到哪里了
                    if cnt % 50 == 0:
                        print(f'{index}:{cnt}')

                    # parse the JSON
                    xdict = json.loads(xline)

                    # Get detail text
                    xdetail = xdict['text']

                    # 解析
                    xs_, xs_type_, xp_, xo_, xo_type_ = None, None, None, None, None
                    spo_list = xdict['spo_list']
                    for xspo in spo_list:
                        # S
                        xs = xspo.get('subject', None)
                        if xs is not None:
                            xs_ = xs
                        xs_type = xspo.get('subject_type', None)
                        if xs_type is not None:
                            xs_type_ = xs_type

                        # P
                        xp = xspo.get('predicate', None)
                        if xp is not None:
                            xp_ = xp

                        # object@value
                        xo = xspo.get('object', None)
                        if xo is not None:
                            xo = xo.get('@value', None)
                            if xo is not None:
                                xo_ = xo

                        # object_type@value
                        xo_type = xspo.get('object_type', None)
                        if xo_type is not None:
                            xo_type = xo_type.get('@value', None)
                            if xo_type is not None:
                                xo_type_ = xo_type

                        # 前面已经检查过，5个值都不是None，所以这里不再判断，直接入库
                        cypher = """\
                            MERGE (s:Node {name: %r, type: %r})
                            MERGE (o:Node {name: %r, type: %r})
                            MERGE (s)-[p:Predicate {name: %r, detail: %r}]->(o)
                            """ % (xs_, xs_type_, xo_, xo_type_, xp_, xdetail)
                        record = ss.run(cypher)
                        print('.', end='')
        print()
        sep('Over')

    def _main():

        # 驱动
        driver = GraphDatabase.driver(**NEO4J_CONFIG_CBLUE_DETAIL)

        # 加索引和约束
        with driver.session() as ss:
            for query in [
                "CREATE INDEX IF NOT EXISTS FOR (n:Node) ON (n.name)",
                "CREATE CONSTRAINT IF NOT EXISTS FOR (n:Node) REQUIRE (n.type, n.name) IS UNIQUE",
                "CREATE INDEX IF NOT EXISTS FOR (p:Predicate) ON (p.name)",
            ]:
                print(query)
                ss.run(query)

        xlist = [
            # r'D:\_dell7590_root\local\LNP_datasets\med\CBLUE\CMeIE-V2\CMeIE-V2_train.jsonl',
            # r'D:\_dell7590_root\local\LNP_datasets\med\CBLUE\CMeIE-V2\CMeIE-V2_dev.limit-20.jsonl',
            '/mnt/d/_dell7590_root/local/LNP_datasets/med/CBLUE/CMeIE-V2/CMeIE-V2_train.jsonl',
            '/mnt/d/_dell7590_root/local/LNP_datasets/med/CBLUE/CMeIE-V2/CMeIE-V2_dev.jsonl',
        ]
        for i, xpath in enumerate(xlist):
            _load(driver, xpath, 20, i)

    _main()
